package com.fy.fyspace.common.limiter;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.TimeUnit;

/**
 * 基于令牌桶算法的限流器实现
 */
@Slf4j
public class TokenBucketRateLimiter implements RateLimiter {
    private final long ratePerSecond;
    private long lastRefillTime;
    private double currentTokens;
    private final long maxTokens;

    public TokenBucketRateLimiter(long ratePerSecond, long maxTokens) {
        this.ratePerSecond = ratePerSecond;
        this.maxTokens = maxTokens;
        this.currentTokens = maxTokens;
        this.lastRefillTime = System.nanoTime();
    }

    @Override
    public synchronized boolean tryAcquire(int tokens) {
        refill();
        if (currentTokens >= tokens) {
            currentTokens -= tokens;
            return true;
        }
        return false;
    }

    @Override
    public synchronized boolean tryAcquire(int tokens, long timeout, TimeUnit unit) {
        long timeoutNanos = unit.toNanos(timeout);
        long startTime = System.nanoTime();
        
        while (System.nanoTime() - startTime < timeoutNanos) {
            if (tryAcquire(tokens)) {
                return true;
            }
            try {
                TimeUnit.MILLISECONDS.sleep(10);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return false;
            }
        }
        return false;
    }

    private void refill() {
        long now = System.nanoTime();
        double tokensToAdd = (now - lastRefillTime) * ratePerSecond / 1_000_000_000.0;
        currentTokens = Math.min(maxTokens, currentTokens + tokensToAdd);
        lastRefillTime = now;
        log.debug("Refilled tokens: current={}, added={}", currentTokens, tokensToAdd);
    }
} 